##########################################################################################

library(data.table)
library(optparse)
library(ggplot2)
library(Seurat)
library(monocle)
library(ArchR)
library(BSgenome.Hsapiens.UCSC.hg38)
library(GenomicFeatures)
library(dplyr)

##########################################################################################
option_list <- list(
    make_option(c("--input_file"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    rna_data_file <- "~/20231121_singleMuti/input/testis_combined.annotationCellType.Rdata"

    ## 单细胞表达文件
    rna_data_monocle_file <- "~/20231121_singleMuti/results/monocole/testis.monocle.Rdata"

    ## atac表达文件
    atac_data_file <- "~/20231121_singleMuti/results/atac_res/testis_combined_peak.Rdata"

    ## 差异表达文件
    time_diff_gene_file <- "~/20231121_singleMuti/results/monocole/pseudotime_ordergene.tsv"

    ## 高变基因的数量
    #variable_num <- 1000

    ## 输出
    out_path <- "~/20231121_singleMuti/results/monocole"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

input_file <- opt$input_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################

a <- load(rna_data_file)
b <- load(atac_data_file)
c <- load(rna_data_monocle_file)

Time_diff <- data.frame(fread(time_diff_gene_file))

###########################################################################################
## RNA的时间
cell_time <- data.frame(cell_id = ALL_cds_all$cell , Pseudotime = ALL_cds_all$Pseudotime )
cell_time$cell_id <- gsub( '_' , '#' , cell_time$cell_id)
## 时间放到100之内,atac需要
cell_time$Pseudotime <- cell_time$Pseudotime/(max(cell_time$Pseudotime)) * 100


###########################################################################################
#### RNA逆时序结果
type <- "testis"

## 黑色的点代码用于构造轨迹的差异基因；灰色是背景基因；
## 红色是根据计算的基因表达大小喝离散度分布的趋势
image_name <- paste0( out_path , "/plot_ordering_genes.",type,".pdf" )
plot2 <- plot_ordering_genes(ALL_cds_all)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)

## 拟时序的时间分布
image_name <- paste0( out_path , "/cell_type_trajectory",type,".pdf" )
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by="cell_type", cell_size=1)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)
image_name <- paste0( out_path , "/cell_type_trajectory",type,".divide.pdf" )
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by = "cell_type") + facet_wrap('~cell_type', nrow = 1)
ggsave(file = image_name , plot = plot2,width = 15,height = 5)

image_name <- paste0( out_path , "/clusters_trajectory",type,".pdf" )
ALL_cds_all$seurat_clusters_character <- as.character(ALL_cds_all$seurat_clusters)
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by="seurat_clusters_character", cell_size=1)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)
image_name <- paste0( out_path , "/clusters_trajectory",type,".divide.pdf" )
ALL_cds_all$seurat_clusters_character <- as.character(ALL_cds_all$seurat_clusters)
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by = "seurat_clusters_character") + facet_wrap('~seurat_clusters_character', nrow = 1)
ggsave(file = image_name , plot = plot2,width = 15,height = 5)

## 时序
image_name <- paste0( out_path , "/pseudotime_trajectory",type,".pdf" )
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by="Pseudotime", cell_size=1)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)


###########################################################################################
## 体现关注基因的染色质激活程度

projHeme5 <- testis_combined_peak
projHeme5@cellColData[cell_time$cell_id,"Pseudotime"] <- cell_time$Pseudotime
#projHeme5 <- addGeneScoreMatrix(projHeme5 , force = T)

## 画轨迹，atac的不太行
p <- plotTrajectory(projHeme5, trajectory = "Pseudotime", colorBy = "cellColData", name = "Pseudotime")
image_name <- paste0( out_path , "/plotTrajectory.atac.pdf" )
pdf(image_name)
p
dev.off()

## 热图
trajGSM <- getTrajectory(ArchRProj = projHeme5, name = "Pseudotime", useMatrix = "GeneScoreMatrix", log2Norm = TRUE)
rownames(trajGSM) <- sapply( strsplit(rownames(trajGSM) , ":") , "[" ,2 )


if(1!=1){
    #trajMM  <- getTrajectory(ArchRProj = projHeme5, name = "Pseudotime", useMatrix = "MotifMatrix", log2Norm = FALSE)
    #trajGSM2 <- trajGSM[corGIM_MM[[1]]$name1, ]
    #trajMM2 <- trajMM[corGIM_MM[[1]]$name2, ]
    #trajCombined <- trajGSM2
    #assay(trajCombined) <- t(apply(assay(trajGSM2), 1, scale)) + t(apply(assay(trajMM2), 1, scale))
    #combinedMat <- plotTrajectoryHeatmap(trajCombined, returnMat = TRUE, varCutOff = 0)
    #p1 <- plotTrajectoryHeatmap(trajMM, pal = paletteContinuous(set = "solarExtra"))
    #corGIM_MM <- correlateTrajectories(trajGSM, trajMM)
}



##  寻找拟时相关的基因（拟时差异基因）
## 把gene放在前面
Time_genes_rna <- (Time_diff %>% pull(gene_short_name) %>% as.character())[1:100]
show_gene <- rownames(trajGSM)[rownames(trajGSM) %in% Time_genes_rna ]

## RNA表达的热图
p <- plot_pseudotime_heatmap(ALL_cds_all[show_gene,] , num_clusters=1 ,show_rownames=T,return_heatmap=T)
image_name <- paste0( out_path , "/plot_pseudotime_heatmap.pdf" )
ggsave(file = image_name , plot = p,width = 15,height = 6)




seTrajectory <- trajGSM
mat <- assay(seTrajectory)




################################################################################
## 构建基因的位置信息

txdb = makeTxDbFromGFF("~/ref/GTF/gencode.v32.annotation.gff3")
gene_ensg <- read.table("~/ref/GTF/gencode.v32.gene_ensg.tsv")

gene_gr <- genes(txdb, columns = c("GENEID"))
colnames(gene_ensg) <- c("GENEID" , "name")

gr_df <- data.frame(
  seqnames = seqnames(gene_gr),
  start = start(gene_gr),
  end = end(gene_gr),
  strand = strand(gene_gr),
  GENEID = names(mcols(gene_gr)$GENEID)
)

## gene name对应多个ensg的，选第一个
gr_df <- merge( gr_df , gene_ensg , by = "GENEID" )
gr_df <- gr_df %>% 
group_by( name ) %>%
summarize( seqnames = seqnames[1] , start = start[1] , end = end[1] , strand = "*"  )

gr <- GRanges(
  seqnames = Rle(gr_df$seqnames),
  ranges = IRanges(start = gr_df$start, end = gr_df$end),
  strand = gr_df$strand,
  name = gr_df$name
)

names(gr) <- gr_df$name

################################################################################

use_gene <- rownames(scrnat@assays$RNA$counts)[rownames(scrnat@assays$RNA$counts) %in% names(gr)]
gr <- gr[use_gene]

seruat_cds_all <- SummarizedExperiment(scrnat@assays$RNA$counts[use_gene,] , rowRanges = gr)
colnames(seruat_cds_all) <- gsub( "_" , "#" ,  colnames(seruat_cds_all) )

## 整合atac和表达，按照每个细胞合并
proj <- addGeneExpressionMatrix(input = testis_combined_peak, seRNA = seruat_cds_all, force = TRUE)
projHeme3 <- addImputeWeights(proj)


## 拟时序
projHeme5 <- projHeme3
projHeme5@cellColData[cell_time$cell_id,"Pseudotime"] <- cell_time$Pseudotime


## atac基因活性的热图
var.cutoff <- 0.9 # Default is 0.9
lims <- c(-1.5, 1.5) # Defaults are c(-1.5, 1.5)
labeltop <- 10 # Default is 50

trajGSM <- getTrajectory(ArchRProj = projHeme5, name = "Pseudotime", useMatrix = "GeneScoreMatrix", log2Norm = TRUE)
trajGIM <- getTrajectory(ArchRProj = projHeme5, name = "Pseudotime", useMatrix = "GeneExpressionMatrix", log2Norm = FALSE)

## Integrative pseudo-time analyses
corGSM_MM <- correlateTrajectories(trajGSM, trajGIM)



p2 <- plotTrajectoryHeatmap(trajGSM, pal=paletteContinuous(set="horizonExtra"), 
    varCutOff=var.cutoff, limits=lims, labelMarkers=show_gene, labelTop=labeltop , maxFeatures = 100000, returnMatrix=FALSE)

p2 <- trajectoryHeatmap(trajGSM[show_gene,],  pal = paletteContinuous(set = "horizonExtra"))
p2 <- plotTrajectoryHeatmap(trajGSM[show_gene,], pal = paletteContinuous(set = "solarExtra"))

image_name <- paste0( out_path , "/plotTrajectoryHeatmap.atac.pdf" )
pdf(image_name)
p
p2
dev.off()

